﻿from simple_pid import PID
import paho.mqtt.client as mqtt
import json
import time
import smbus2,ctypes
import Hobot.GPIO as GPIO
import atexit  # 用于在程序退出时自动调用某些函数

class mpu_object:

    def __init__(self):
        self.bus = smbus2.SMBus(0)
        #固定参数
        self.MPU6050_ADDR = 0x68
        self.PWR_MGMT_1   = 0x6B
        self.SMPLRT_DIV   = 0x19
        self.CONFIG       = 0x1A
        self.GYRO_CONFIG  = 0x1B
        self.ACCEL_CONFIG = 0x1C
        self.INT_ENABLE   = 0x38
        self.ACCEL_XOUT_H = 0x3B
        self.ACCEL_YOUT_H = 0x3D
        self.ACCEL_ZOUT_H = 0x3F
        self.TEMP_OUT_H   = 0x41
        self.GYRO_XOUT_H  = 0x43
        self.GYRO_YOUT_H  = 0x45
        self.GYRO_ZOUT_H  = 0x47

        samp_rate_div = 7 # sample rate = 8 kHz/(1+samp_rate_div)
        self.bus.write_byte_data(self.MPU6050_ADDR,self.SMPLRT_DIV, samp_rate_div)

        self.bus.write_byte_data( self.MPU6050_ADDR, self.PWR_MGMT_1,0x00)
        time.sleep(0.1)
        self.bus.write_byte_data(self.MPU6050_ADDR, self.PWR_MGMT_1, 0x01)
        time.sleep(0.1)
        self.bus.write_byte_data( self.MPU6050_ADDR,  self.CONFIG, 0)
        time.sleep(0.1)

        #Write to Gyro configuration register
        gyro_config_sel = [0b00000,0b010000,0b10000,0b11000] # byte registers
        gyro_config_vals = [250.0,500.0,1000.0,2000.0] # degrees/sec
        gyro_indx = 0
        self.bus.write_byte_data( self.MPU6050_ADDR,  self.GYRO_CONFIG, int(gyro_config_sel[gyro_indx]))
        time.sleep(0.1)
        #Write to Accel configuration register
        accel_config_sel = [0b00000,0b01000,0b10000,0b11000] # byte registers
        accel_config_vals = [2.0,4.0,8.0,16.0] # g (g = 9.81 m/s^2)
        accel_indx = 0
        self.bus.write_byte_data( self.MPU6050_ADDR,  self.ACCEL_CONFIG, int(accel_config_sel[accel_indx]))
        time.sleep(0.1)
        self.gyro_sens = gyro_config_vals[gyro_indx]
        self.accel_sens=accel_config_vals[accel_indx]


        # x_accel_offset = int(-1588)
        # y_accel_offset = int(1520.875)
        # z_accel_offset = int(-755.5)
        # x_gyro_offset = int(16.5)
        # y_gyro_offset = int(-13.25)
        # z_gyro_offset = int(-25)
        # 
        # if x_accel_offset != 0:
        #     self.set_x_accel_offset(x_accel_offset)
        # 
        # if y_accel_offset != 0:
        #     self.set_y_accel_offset(y_accel_offset)
        # 
        # if z_accel_offset != 0:
        #     self.set_z_accel_offset(z_accel_offset)
        # 
        # if x_gyro_offset != 0:
        #     self.set_x_gyro_offset(x_gyro_offset)
        # 
        # if y_gyro_offset != 0:
        #     self.set_y_gyro_offset(y_gyro_offset)
        # 
        # if z_gyro_offset != 0:
        #     self.set_z_gyro_offset(z_gyro_offset)

    def read_raw_bits(self,register):
        # read accel and gyro values
        high = self.bus.read_byte_data(self.MPU6050_ADDR, register)
        low = self.bus.read_byte_data(self.MPU6050_ADDR, register+1)

        # combine higha and low for unsigned bit value
        value = ((high << 8) | low)

        # convert to +- value
        if(value > 32768):
            value -= 65536
        return value

    def read_accel(self):
        # read raw accel values
        x = self.read_raw_bits(self.ACCEL_XOUT_H)
        y = self.read_raw_bits(self.ACCEL_YOUT_H)
        z = self.read_raw_bits(self.ACCEL_ZOUT_H)

        scale = 32768/self.accel_sens

        x = x / scale
        y = y / scale
        z = z / scale
        return x,y,z

    def read_gyro(self):
        # read raw gyro values
        x = self.read_raw_bits(self.GYRO_XOUT_H)
        y = self.read_raw_bits(self.GYRO_YOUT_H)
        z = self.read_raw_bits(self.GYRO_ZOUT_H)

        scale = 32768/self.gyro_sens

        x = x / scale
        y = y / scale
        z = z / scale
        return x,y,z

    def set_x_accel_offset(self,a_offset):
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x06,ctypes.c_int8(a_offset >> 8).value)
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x07,ctypes.c_int8(a_offset).value)

    def set_y_accel_offset(self,a_offset):
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x08,ctypes.c_int8(a_offset >> 8).value)
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x09,ctypes.c_int8(a_offset).value)

    def set_z_accel_offset(self,a_offset):
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x0A,ctypes.c_int8(a_offset >> 8).value)
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x0B,ctypes.c_int8(a_offset).value)

    def set_x_gyro_offset(self,g_offset):
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x13,ctypes.c_int8(g_offset >> 8).value)
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x14,ctypes.c_int8(g_offset).value)

    def set_y_gyro_offset(self,g_offset):
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x15,ctypes.c_int8(g_offset >> 8).value)
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x16,ctypes.c_int8(g_offset).value)

    def set_z_gyro_offset(self,g_offset):
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x17,ctypes.c_int8(g_offset >> 8).value)
        self.bus.write_byte_data(self.MPU6050_ADDR, 0x18,ctypes.c_int8(g_offset).value)


#gpio设定

# 清理可能已经被导出的通道
GPIO.cleanup()

# GPIO初始化
GPIO.setmode(GPIO.BOARD)

# 接口编号
LEFT_PWM_PIN = 33   # 左轮前向PWM接口
LEFT_GPIO_PIN = 37  # 左轮后向GPIO接口
RIGHT_PWM_PIN = 32  # 右轮前向PWM接口
RIGHT_GPIO_PIN = 36 # 右轮后向GPIO接口

# 初始化GPIO
GPIO.setup(LEFT_GPIO_PIN, GPIO.OUT)
GPIO.setup(RIGHT_GPIO_PIN, GPIO.OUT)

# 初始化PWM信号
left_p_obj = GPIO.PWM(LEFT_PWM_PIN, 48000)  # PWM频率48kHz
right_p_obj = GPIO.PWM(RIGHT_PWM_PIN, 48000)

# 设置初始状态
left_p_obj.ChangeDutyCycle(1)
right_p_obj.ChangeDutyCycle(1)
left_p_obj.start(0)
right_p_obj.start(0)


# MQTT配置
broker_address = "127.0.0.1"  # 替换为实际的MQTT Broker地址
topic = "sensor/data"  # 要发布的主题
# 创建MQTT客户端并连接到Broker
client = mqtt.Client()
client.connect(broker_address, 1883, 60)

def mqtt_set_msg(target,pid_data):
    payload = {
        "target":target,
        "control":pid_data,
    }
    payload_json = json.dumps(payload)
    client.publish(topic, payload_json)

def read_gyro_z():
    # a_x,a_y,a_z = mpu.read_accel()
    g_x,g_y,g_z = mpu.read_gyro()
    
    # if(abs(g_z)<0):
    #     return 0
    
    return g_z

def set_motor_pwm_gpio(left_pwm, left_gpio, right_pwm, right_gpio):
    # 这里设置左右电机的PWM和GPIO值
    payload = {
        "left_pwm":left_pwm,
        "right_pwm":right_pwm,
        # "left_gpio":left_gpio,
        # "right_gpio":right_gpio,
    }
    payload_json = json.dumps(payload)
    client.publish(topic, payload_json)
    
    print(left_pwm, left_gpio, right_pwm, right_gpio)
    
    #gpio逻辑
    left_p_obj.ChangeDutyCycle(int(left_pwm))
    GPIO.output(LEFT_GPIO_PIN, left_gpio)
    right_p_obj.ChangeDutyCycle(int(right_pwm))
    GPIO.output(RIGHT_GPIO_PIN, right_gpio)

def set_motor_group_value(fix_group, dynamic_group,base_pwm,base_gpio,dynamic_pwm):
    
    fix_group["pwm"].ChangeDutyCycle(int(base_pwm))
    GPIO.output(fix_group["gpio"], base_gpio)
    dynamic_group["pwm"].ChangeDutyCycle(int(dynamic_pwm))
    GPIO.output(dynamic_group["gpio"], base_gpio)

    
    print(base_pwm, dynamic_pwm,base_gpio)


    payload = {
        "fixed_pwm":base_pwm,
        "dynamic_pwm":dynamic_pwm
      
    }

    payload_json = json.dumps(payload)
    client.publish(topic, payload_json)


mpu = mpu_object()

def clamp(value, min_value, max_value):
    return max(min_value, min(value, max_value))

# @atexit.register
def cleanup():
    print("Cleaning up resources...")
    client.loop_stop()  # 停止 MQTT 监听
    left_p_obj.ChangeDutyCycle(1)
    right_p_obj.ChangeDutyCycle(1)
    left_p_obj.start(0)
    right_p_obj.start(0)

    GPIO.output(LEFT_GPIO_PIN, 0)
    GPIO.output(RIGHT_GPIO_PIN, 0)
    
    GPIO.cleanup([LEFT_PWM_PIN,LEFT_GPIO_PIN,RIGHT_PWM_PIN,RIGHT_GPIO_PIN])

# 卡尔曼滤波器类
class KalmanFilter:
    def __init__(self, Q=1e-5, R=0.01):
        self.Q = Q  # 过程噪声协方差
        self.R = R  # 测量噪声协方差
        self.x = 0.0  # 初始估计值
        self.P = 1.0  # 初始估计协方差
        self.K = 0.0  # 初始卡尔曼增益

    def update(self, measurement):
        # 预测步骤
        self.P = self.P + self.Q  # 更新估计协方差
        # 计算卡尔曼增益
        self.K = self.P / (self.P + self.R)
        # 更新估计值
        self.x = self.x + self.K * (measurement - self.x)
        # 更新协方差
        self.P = (1 - self.K) * self.P
        return self.x
    
    def clean(self, Q=1e-5, R=0.01):
        self.Q = Q  # 过程噪声协方差
        self.R = R  # 测量噪声协方差
        self.x = 0.0  # 初始估计值
        self.P = 1.0  # 初始估计协方差
        self.K = 0.0  # 初始卡尔曼增益
    
kf = KalmanFilter()

dt = 0.1


#假如一侧比较快 一侧比较慢  那么慢的就是基准不变，调节快的电机

def analyze_single_container(data):
    right_turns = 0
    left_turns = 0

    # 计算容器中的差值并分析偏向趋势
    for i in range(1, len(data)):
        difference = data[i] - data[i-1]
        if difference > 0:
            right_turns += 1
        elif difference < 0:
            left_turns += 1

    # 根据正负差值的数量确定总体趋势
    if right_turns > left_turns:
        return "Right"
    elif left_turns > right_turns:
        return "Left"
    else:
        return "Stable"
    

base_mark = "Left"

def determine_base_motor():
    left_pwm = 100
    left_gpio = 0
    right_pwm = 100
    right_gpio = 0
    set_motor_pwm_gpio(left_pwm, left_gpio, right_pwm, right_gpio)
    det_time = time.time()
    
    f_data = []
    
    while time.time() - det_time < 1:
        #统计偏航数据
        gyro_z = read_gyro_z()
        f_data.append(gyro_z)
        time.sleep(0.1)

    left_pwm = 0
    left_gpio = 1
    right_pwm = 0
    right_gpio = 1
    set_motor_pwm_gpio(left_pwm, left_gpio, right_pwm, right_gpio)
    det_time = time.time()

    d_data = []
    
    while time.time() - det_time < 1:
        #统计偏航数据
        gyro_z = read_gyro_z()
        d_data.append(gyro_z)
        time.sleep(0.1)
    
    # 判定趋势基础
    # 分别计算 f_data 和 d_data 的偏向结果
    direction_f_data = analyze_single_container(f_data)
    direction_d_data = analyze_single_container(d_data)

    if direction_f_data == direction_d_data:
        base_mark = direction_f_data
    else:
        base_mark = ""


    cleanup()

# determine_base_motor()

def distribute_motor_group():

    if base_mark == "Left":
        #标记基础驱动电机
        fixed_motor_group = {
            "pwm": left_p_obj,
            "gpio":LEFT_GPIO_PIN,
            "dir":1 
        }
        
        #标记动态驱动电机
        dynamic_motor_group = {
            "pwm": right_p_obj,
            "gpio":RIGHT_GPIO_PIN,
            "dir":-1
        }
        return fixed_motor_group,dynamic_motor_group
    
    if base_mark == "Right":
        #标记基础驱动电机
        fixed_motor_group = {
            "pwm": right_p_obj,
            "gpio":RIGHT_GPIO_PIN,
            "dir":1
        }
        
        #标记动态驱动电机
        dynamic_motor_group = {
            "pwm": left_p_obj,
            "gpio":LEFT_GPIO_PIN,
            "dir":-1
        }
        return fixed_motor_group,dynamic_motor_group
    
    return None
    

    

if __name__ == '__main__':

    fixed_group,dynamic_group = distribute_motor_group()
    #初始化控制pid
    pid_dynamic_group = PID(10, 1, 1, setpoint=0)
    # pid_dynamic_group = PID(1, 0, 0.1, setpoint=0)
    pid_dynamic_group.output_limits = (-180,180)

    # 初始化两个PID控制器
    # pid_forward = PID(1, 0.1, 0.05, setpoint=0)  # 前进时的PID
    # pid_backward = PID(1, 0.1, 0.05, setpoint=0)  # 后退时的PID
    # 
    # # 设置PID的输出限制，控制PWM调整范围
    # pid_forward.output_limits = (-180,180)
    # pid_backward.output_limits = (-180,180)

    #前进的基准值
    base_f_pwm = 100
    #后退的基准值
    base_b_pwm = 0
    direction = 'backward'  # 初始状态为前进，可以设置为'backward'表示后退

    cur_dynamic_pwm = 0

    
    if direction == 'forward':
        base_dynamic_pwm = base_f_pwm
       
    elif direction == 'backward':
        base_dynamic_pwm = base_b_pwm
       
    else:
        raise ValueError("Invalid direction value. Please set it to 'forward' or 'backward'.")

    cur_dynamic_pwm = base_dynamic_pwm

    stable_value = 0
    stable_time = time.time()
    stable_yaw = 0


    f_mark = False

    #稳定值
    while time.time() - stable_time < 1.5:
        stable_value = kf.update(read_gyro_z())
        stable_yaw = 0

    try:

        while True:
            gyro_z = read_gyro_z()
            if direction == 'forward':

                if abs(gyro_z - stable_value) < 0.5:
                    gyro_z = 0
                else :
                    gyro_z -= stable_value

                stable_yaw += gyro_z*dt
                correction = pid_dynamic_group(stable_yaw)

                #增量

                correction = int(correction)
                print(correction)

                #动态侧调整
                cur_dynamic_pwm = clamp(cur_dynamic_pwm + correction,50,100)
             
                fixed_gpio = 0
                dynamic_pwm = cur_dynamic_pwm;
                

            elif  direction == 'backward':
                # correction = pid_backward(gyro_z)

                # 后退状态，调整PWM值
                # left_pwm = base_b_pwm + correction
                # right_pwm = base_b_pwm - correction
                # left_gpio = 1
                # right_gpio = 1
                # 
                # 
                # if abs(gyro_z - stable_value) < 1:
                #     gyro_z = 0
                # else :
                #     print("gyro_z:",gyro_z-stable_value)
                #     gyro_z -= stable_value
                # 
                # stable_yaw += gyro_z*dt
                # correction = pid_backward(stable_yaw)
                # 
                # #增量
                # correction = int(correction)
                # 
                # print(correction)
                # 
                # # 后退状态，调整PWM值
                # 
                # cur_left_pwm = clamp(cur_left_pwm - correction,0,30)
                # cur_right_pwm = clamp(cur_left_pwm + correction,0,30)
                # 
                # left_pwm = cur_left_pwm
                # right_pwm = cur_right_pwm

                left_pwm = 0
                right_pwm = 100

                if abs(gyro_z - stable_value) < 0.5:
                    gyro_z = 0
                else :
                    gyro_z -= stable_value

                stable_yaw += gyro_z*dt
                correction = pid_dynamic_group(stable_yaw)

                #增量

                correction = int(correction)
                print(correction)

                #动态侧调整
                cur_dynamic_pwm = clamp(cur_dynamic_pwm + correction,0,50)

                
                
                fixed_gpio = 1
                dynamic_pwm = cur_dynamic_pwm;



            # set_motor_pwm_gpio(left_pwm, left_gpio, right_pwm, right_gpio)

            set_motor_group_value(fixed_group,dynamic_group,base_dynamic_pwm,fixed_gpio,dynamic_pwm)


            time.sleep(dt)
    except KeyboardInterrupt:
        cleanup()

